# -*- coding: utf-8 -*-
"""

数值处理类

def grid_interp(); 插值
def grid_discrete(); 差分/微分
                
"""
# base lib
import sys

# third party lib
import numpy as np
from scipy import linalg
from scipy.signal import savgol_filter

# private lib
from utils.constant import InterpMode, DifferenceMode, ExtrapMode


def grid_interp(data: np.ndarray, mode: int, velocity=0) -> np.ndarray:
    """插值函数

    Args:
        data: numpy ndarray; 变量矩阵
        mode: int; 差分方式：
                1.'mid'; 中心插值
                2.'fou'; 迎风格式
        velocity: int

    Returns:
        np.ndarray,插值结果
    """
    if mode == InterpMode.MID.value:
        result = (data[1:] + data[:-1]) / 2
    elif mode == InterpMode.FOU.value:
        result = np.where(velocity > 0, data[:-1], data[1:])
    return result


def grid_extrap(data: np.ndarray, mode: int) -> np.ndarray:
    """外插函数

    Args:
        data: numpy ndarray; 变量矩阵
        mode: int; 差分方式：
                1.'mid in + linear boundary'; 中心插值 + 线性外推边界
                2. ... 待添加

    Returns:
        np.ndarray,插值结果
    """
    if mode == ExtrapMode.MIDin_linearB.value:
        result = np.zeros(len(data) + 1)
        result[1:-1] = grid_interp(data, 1)  # 内部的值采用中心差分
        result[0] = 2 * data[0] - data[1]  # 外部的值采用线性外推
        result[-1] = 2 * data[-1] - data[-2]
    return result


def difference(data: np.ndarray, dx: float, mode: int) -> np.ndarray:
    """差分/微分:

    Args:
        data: numpy ndarray; 变量矩阵
        dx: float; 微分最小单元/网格物理长度
        mode: int; 差分方式：
              1. 单步长中心差分格式
              2. 双步长中心差分格式
              3. 迎风格式

    Returns:
        np.ndarray,差分/微分结果
    """

    if mode == DifferenceMode.CDMs.value:
        result = (data[1:] - data[:-1]) / dx
        return result
    if mode == DifferenceMode.CDMd.value:
        result = (data[2:] - data[:-2]) / 2 / dx
        return result
    if mode == DifferenceMode.FOU.value:
        pre_gradient = data[1:-1] - data[:-2]
        post_gradient = data[2:] - data[1:-1]
        return np.where(data[1:-1] > 0, pre_gradient, post_gradient) / dx
    # if mode == DifferenceMode.MinMod.value:  # MinMOD
    #     var_m = data[1:-1] + 0.5 * minmod(data[1:-1] - data[:-2], data[2:] - data[1:-1])
    #     var_left = data[:-2] + 0.5 * minmod(
    #         data[:-2] - np.insert(data, 0, data[0])[:-3], data[1:-1] - data[:-2]
    #     )
    #     pre_gradient = var_m - var_left

    #     var_right = data[2:] + 0.5 * minmod(
    #         data[2:] - data[1:-1], np.insert(data, -1, data[-1])[3:] - data[2:]
    #     )
    #     post_gradient = var_right - var_m
    #     return np.where(data[1:-1] > 0, pre_gradient, post_gradient) / dx
    sys.exit("请确认输入模式的正确性")

def penta_diagonal_solve(a_P, a_E, a_W, a_EE, a_WW, b):
    """使用scipy.linalg.solve_banded求解五对角矩阵方程 Ax = b"""
    n = len(a_P)
    
    # 1. 检查对角占优并调整
    diag_sum = np.abs(a_P)
    off_diag = np.zeros_like(a_P)
    off_diag[1:] += np.abs(a_W[1:])    
    off_diag[2:] += np.abs(a_WW[2:])   
    off_diag[:-1] += np.abs(a_E[:-1])  
    off_diag[:-2] += np.abs(a_EE[:-2])
    
    # 调整对角元确保对角占优
    mask = diag_sum <= off_diag
    if np.any(mask):
        # log.warning("矩阵可能病态")
        # log.info(f"mask: {mask}, off_diag: {off_diag[mask]}, diag_sum: {diag_sum[mask]}")
        raise RuntimeError("矩阵可能病态")
        # a_P = np.where(mask, off_diag * 1.1, a_P)
        
    # 2. 添加小扰动避免奇异
    eps = 1e-10
    a_P += eps * np.ones_like(a_P)
    
    # 3. 构造带状矩阵
    ab = np.zeros((5, n))
    ab[0, 2:] = a_EE[:-2]  
    ab[1, 1:] = a_E[:-1]   
    ab[2, :] = a_P        
    ab[3, :-1] = a_W[1:]   
    ab[4, :-2] = a_WW[2:]  
    try:
        x = linalg.solve_banded((2, 2), ab, b)
    except np.linalg.LinAlgError:
        # 如果求解失败,使用更稳定的求解器
        raise RuntimeError("求解失败")
    return x

def thomas_algorithm(a_W, a_P, a_E, b):
    """ 使用向量化的Thomas算法求解三对角矩阵方程 Ax = b
    
    Args:
        a_W: 下对角元素 [a2, a3, ..., an]
        a_P: 主对角元素 [b1, b2, ..., bn]
        a_E: 上对角元素 [c1, c2, ..., c(n-1)]
        b: 右端项 [d1, d2, ..., dn]
    
    Returns:
        x: 方程的解
    """
    n = len(b)
    
    # 1. 检查对角占优性并调整
    # 构造非对角元素和
    off_diag = np.zeros_like(a_P)
    off_diag[1:] += np.abs(a_W[1:])    # 下对角
    off_diag[:-1] += np.abs(a_E[:-1])  # 上对角
    
    # 检查并调整对角占优
    diag_sum = np.abs(a_P)
    a_P = np.where(diag_sum <= off_diag, off_diag * 1.1, a_P)

    # 2. 构造scipy需要的带状矩阵格式
    ab = np.zeros((3, n))
    
    # 填充带状矩阵
    ab[0, 1:] = a_E[:-1]   # 上对角
    ab[1, :] = a_P         # 主对角
    ab[2, :-1] = a_W[1:]   # 下对角
    
    # 3. 使用scipy的带状矩阵求解器
    x = linalg.solve_banded((1, 1), ab, b)
    
    return x

def bisection_method(func, up_bound, low_bound, tol=1e-6):
    if func(up_bound) * func(low_bound) >= 0:
        raise RuntimeError("二分法求解失败, 该方程无解")
    
    while (up_bound - low_bound) / 2 > tol:
        center = (up_bound + low_bound) / 2
        if func(center) == 0:
            return center
        elif func(up_bound) * func(center) < 0:
            up_bound = center
        else:
            low_bound = center
    
    return (up_bound + low_bound) / 2


def gaussian_smooth(data: np.ndarray, sigma: float) -> np.ndarray:
    """ 高斯平滑
    
    Args:
        data: np.ndarray; 数据
        sigma: float; 标准差
    
    Returns:
        np.ndarray, 平滑后的数据
    """
    window_size = int(4 * sigma)  # 窗口大小设为6倍标准差
    if window_size % 2 == 0:
        window_size += 1  # 确保窗口大小为奇数
    x = np.arange(-window_size // 2 + 1, window_size // 2 + 1)
    gaussian_kernel = np.exp(-x ** 2 / (2 * sigma ** 2))
    gaussian_kernel = gaussian_kernel / gaussian_kernel.sum()
    return np.convolve(data, gaussian_kernel, mode='same')


def simple_moving_average(data, window_size):
    """简单移动平均
    参数:
        data: 输入数据
        window_size: 窗口大小
    """
    weights = np.ones(window_size) / window_size
    sma = np.convolve(data, weights, mode='valid')
    
    # 处理边界，保持数据长度一致
    pad_size = window_size - 1
    pad_left = np.full(pad_size//2, sma[0])
    pad_right = np.full(pad_size - pad_size//2, sma[-1])
    return np.concatenate((pad_left, sma, pad_right))


def apply_global_linearization(data_series, weight=0.3):
    """
    全局线性约束算法
    :param data_series: 原始数据序列 (numpy数组)
    :param weight: 线性约束强度 (0~1)
    :return: 平滑后的序列
    """
    n = len(data_series)
    X = np.vstack([np.ones(n), np.arange(n)]).T
    
    # 构造混合矩阵
    A = weight * np.eye(n) + (1-weight) * X @ np.linalg.inv(X.T @ X) @ X.T
    return A @ data_series

def apply_global_quadratic_constraint(data_series, weight=0.3):
    """
    全局二次约束算法
    :param data_series: 原始数据序列 (numpy数组)
    :param weight: 约束强度 (0~1)
    :return: 约束后的序列
    """
    n = len(data_series)
    x = np.arange(n)
    # 构造二次回归设计矩阵 [1, x, x²]
    X = np.vstack([np.ones(n), x, x**2, x**3, x**4, x**5]).T

    # 计算投影矩阵 H = X(X^T X)^-1 X^T
    H = X @ np.linalg.inv(X.T @ X) @ X.T
    # 构造混合矩阵
    A = weight * np.eye(n) + (1 - weight) * H
    
    return A @ data_series